package auth

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"kratos-layout/pkg/jwt/claim"

	"github.com/go-kratos/kratos/v2/transport"
	"github.com/golang-jwt/jwt/v4"
)

const (
	UserIDKey        = "UserID"
	AuthorizationKey = "Authorization"
	BearerKey        = "Bearer"
)

var (
	ErrUnauthorized = fmt.Errorf("Unauthorized")
	ErrJwtExpired   = fmt.Errorf("token is expired")
	ErrTokenFormat  = fmt.Errorf("token format error")
	ErrTokenMissing = fmt.Errorf("token is missing")
)

type JWTTokenVerifier struct {
	secretKey []byte
	refresh   bool //是否支持续期
}

func NewJWTTokenVerifier(secretKey string, refresh bool) *JWTTokenVerifier {
	return &JWTTokenVerifier{
		secretKey: []byte(secretKey),
		refresh:   refresh,
	}
}

func (j *JWTTokenVerifier) Verify(tkn string) (int, error) {
	t, err := jwt.ParseWithClaims(tkn, &claim.CustomClaims{}, func(token *jwt.Token) (interface{}, error) {
		return j.secretKey, nil
	})
	if err != nil {
		if ve, ok := err.(*jwt.ValidationError); ok {
			if ve.Errors&jwt.ValidationErrorMalformed != 0 {
				return 0, fmt.Errorf("token malformed: %w", err)
			} else if ve.Errors&jwt.ValidationErrorExpired != 0 {
				if j.refresh {
					claims, ok := t.Claims.(*claim.CustomClaims)
					if ok {
						return claims.UserId, ErrJwtExpired
					}
				}
				return 0, fmt.Errorf("token is expired: %w", err)
			} else if ve.Errors&jwt.ValidationErrorNotValidYet != 0 {
				return 0, fmt.Errorf("token is not valid yet: %w", err)
			} else {
				return 0, fmt.Errorf("token validation error: %w", err)
			}
		}
		return 0, fmt.Errorf("cannot parse token: %w", err)
	}
	if !t.Valid {
		return 0, fmt.Errorf("token not valid")
	}
	claims, ok := t.Claims.(*claim.CustomClaims)
	if !ok {
		return 0, fmt.Errorf("claim type error")
	}
	if err := claims.Valid(); err != nil {
		return 0, fmt.Errorf("claims not valid")
	}
	return claims.UserId, nil
}

func GetToken(ctx context.Context) (string, error) {
	if header, ok := transport.FromServerContext(ctx); ok {
		auths := strings.SplitN(header.RequestHeader().Get(AuthorizationKey), " ", 2)
		if len(auths) != 2 || !strings.EqualFold(auths[0], BearerKey) {
			return "", ErrTokenFormat
		}
		return auths[1], nil
	}
	return "", ErrTokenMissing
}

func VerifyToken(token string, secretKey string) (int, error) {
	if token == "" {
		return 0, ErrUnauthorized
	}
	token = strings.TrimPrefix(token, BearerKey)
	token = strings.TrimSpace(token)
	verifier := NewJWTTokenVerifier(secretKey, true)
	userID, err := verifier.Verify(token)
	if err != nil {
		if errors.Is(err, ErrJwtExpired) {
			return userID, err
		}
		return 0, err
	}
	if userID == 0 {
		return 0, ErrUnauthorized
	}
	return userID, nil
}
